# encoding: utf8


import h5py as h5
import tables as tb5
import os
import numpy as np
import pandas as pd
import time
import matplotlib.pyplot as plt


class DemoHdf5:
    """
    生成一些图像数据，
    保存二进制图像数据到hdf5，并读出为图像数据
    二进制数据格式转换
    测试存储、读取、转换效率
    """

    def __init__(self, testfile='syntax.hdf5'):
        self.h5file = testfile
        self.image = None

    def test_write(self, group='group', num=10):
        st = time.time()
        mode = 'w' if not os.path.isfile(self.h5file) else 'r+'
        print(mode)
        with h5.File(self.h5file, mode=mode) as fh:
            # 需要时重新开创组 group
            fh.require_group(group)
            # 模拟读取多个图像写入组 group
            # 写入图像数据
            images = ['xin.png', 'textdoc.jpg', 'textdoc2.jpg', 'textdoc3.jpg']
            for j in range(num):
                expimage = images[j % 4]
                with open(expimage, 'rb') as fp:
                    bs = fp.read()
                    image = np.array(bytearray(bs), dtype=np.uint8)
                    image_name = group + '/{:05d}--{}'.format(j, expimage)
                    if image_name.split('/')[-1] in fh[group].keys():
                        dset = fh[group]
                        dset = image
                    else:
                        fh[image_name] = image
                if j % 500 == 0:
                    print('written images {:04d}, len={}'.format(j, len(image)))
        print(time.time()-st)

    def test_read(self, disp=False):
        # 读取组group中的图像数据
        # 测试读取时间:
        # ever: 412 with print each name
        #       454 for 1/10000, 94 read index
        # best: 271 for print 1/10000 with 28 to read index

        st = time.time()
        read_index = False
        with h5.File(self.h5file) as hfp:
            # 遍历组 group
            for imgname in hfp['group'].keys():
                image = bytes(np.array(hfp['group/'+imgname], dtype=np.uint8))
                if not read_index:
                    read_index = True
                    print(time.time()-st)
                if int(imgname[imgname.find('--')+2:]) % 10000 == 0:
                    print('{}.len == {}'.format(imgname, len(image)))
        print(time.time()-st)

        # 转为文件
        # with open('testtemp.png', 'wb') as fp:
        #     fp.write(image)
        # img = plt.imread('testtemp.png')

        # 直接转换
        import cv2
        img = cv2.imdecode(np.asarray(bytearray(image), dtype='uint8'), cv2.IMREAD_COLOR)

        # 验证显示
        if disp:
            plt.imshow(img)
            self.image = image

    def test_time_numpy_read_name(self, generator=False):
        st = time.time()
        with h5.File(self.h5file) as fh5:
            namelist = list(fh5['group'].keys()) if not generator else fh5['group'].keys()
        print(len(namelist))
        print(time.time()-st)

    def test_pandas(self):
        pass
        # df = pd.DataFrame({'image'+str(j): [v for v in bs]})
        # df.to_hdf('syntax.hdf5', 'group1/image{:02d}'.format(j))
        # df0 = pd.read_hdf('syntax.hdf5', 'image_group/image00')
        # bs0 = bytes([int(v) for v in df0.image0])
        # print(bs0)


def read_bin_data_to_image(read_image_path, write_image_path):
    import cv2

    # 解码二进制流
    with open(read_image_path, "rb") as file:
        jpg_bin = file.read()
        image = cv2.imdecode(np.asarray(bytearray(jpg_bin), dtype='uint8'), cv2.IMREAD_COLOR)

    # 编码Jpg存储数据
    with open(write_image_path, 'wb') as tmp_file:
        tmp_jpg_bin = np.array(cv2.imencode('.jpg', image)[1]).tobytes()
        tmp_file.write(tmp_jpg_bin)


class Hdf5:

    def __init__(self, h5file='demo.hdf5'):
        self.fh5 = h5file

    def new_group(self, group='group01'):
        with h5.File(self.fh5, 'w') as fh:
            fh.require_group(group)

    def add_images(self, group: str, images: dict):
        with h5.File(self.fh5, 'w') as fh:
            fh.require_group(group)
            for dataname in images.keys():
                data = np.array(bytearray(images[dataname]), dtype=np.uint8)
                fh[group + '/' + dataname] = data

    def add_images_from_path(self, group, path, suffix=('jpg', 'png')):
        """
        将图像文件从目录path中读出，存入组group中
        约定：数据集名使用文件名
        :param group: 写入图像文件数据的组名
        :param path: 图像文件存放路径
        :param suffix: 图像文件后缀，非列表中后缀名的文件排除
        """
        path = path.replace('\\', '/')
        suffix_upper = [s.upper() for s in suffix]
        with h5.File(self.fh5, 'w') as fh5:
            fh5.require_group(group)
            for file in os.listdir(path):
                if '.' not in file:
                    continue
                if file.split('.')[-1].upper() not in suffix_upper:
                    continue
                print('write image file: {}'.format(file))
                with open(file, 'rb') as fp:
                    image_name = os.path.split(file)[-1]
                    fh5[group + '/' + image_name] = np.array(bytearray(fp.read()), dtype=np.uint8)

    def read_iamge(self, group, path='temp_images'):
        """
        将图像数据从group读出，写入目录Path之中
        图像文件名使用数据集名
        约定：存入数据时，数据集名使用文件名
        要点：使用bytes(image[:len(image)])转换为二进制字节数据，时间大大降低
        """
        if not os.path.isdir(path):
            os.mkdir(path)
        with h5.File(self.fh5, 'r') as fh:
            print('read data name list...')
            names = fh['group'].keys()
            print('read data ...')
            for ds in names:
                image = fh[group+'/'+ds]
                with open(path+'/ds', 'wb') as fp:
                    fp.write(bytes(image[:len(image)]))

    def info(self):
        files_count = 0
        with h5.File(self.fh5, 'r') as fh:
            for gf in fh.keys():
                if isinstance(fh[gf], h5.Group):
                    print(gf)
                    self.__get_group_files(fh, gf, 1)
                else:
                    files_count += 1
            print('/ files_count={}'.format(files_count))

    def __get_group_files(self, fh, group, level=0):
        file_count = 0
        for gf in fh[group].keys():
            if isinstance(fh[group+'/'+gf], h5.Group):
                print(' '*level*3 + group+'/'+gf)
                self.__get_group_files(fh, group+'/'+gf, level=level+1)
            else:
                file_count += 1
        print(' '*level*3+group+'/'+' files_count={}'.format(file_count))
